# -*- coding: UTF-8 -*-
# -------------------------------------------------------------------------------
# Filename:     test-cnmon-log2png.py
# UpdateDate:   2024/06/08
# Description:  一键导出test-cnmon测试日志中关键信息并生成图片。
# Example:      python test-cnmon_log2png.py ./log/test-cnmon-20240520143052.601923295.log ./test.png
# Depends:      如没安装matplotlib库，可以通过pip安装： pip install matplotlib
# Notes:        本脚本适用于Docker容器环境下运行；
# -------------------------------------------------------------------------------
import re
import sys
import matplotlib.pyplot as plt

#def plot_system_metrics(log_file_path, output_image_path):
def plot_system_metrics(log_file_path, output_image_path='default_system_metrics.png'):
    # 用于存储每个关键字对应数值的列表
    ipu_frequencies = []
    mlu_averages = []
    usages = []

    # 打开日志文件并读取每一行
    with open(log_file_path, 'r') as file:
        for line in file:
            # 提取IPU频率
            ipu_match = re.search(r'IPU\s*:\s*(\d+) MHz', line)
            if ipu_match:
                ipu_frequencies.append(int(ipu_match.group(1)))

            # 提取MLU Average百分比
            mlu_match = re.search(r'MLU Average\s*:\s*(\d+) %', line)
            if mlu_match:
                mlu_averages.append(int(mlu_match.group(1)))

            # 提取Usage功率
            usage_match = re.search(r'Usage\s*:\s*(\d+) W', line)
            if usage_match:
                usages.append(int(usage_match.group(1)))

    # 假设我们简单地使用索引作为x轴的值
    x = range(len(ipu_frequencies))

    # 绘制三条曲线图
    plt.figure(figsize=(10, 8))

    plt.subplot(3, 1, 1)
    plt.plot(x, ipu_frequencies, marker='o', label='IPU Frequency (MHz)')
    plt.title('IPU Frequency Over Time')
    plt.xlabel('Index')
    plt.ylabel('IPU Frequency (MHz)')
    plt.grid(True)

    plt.subplot(3, 1, 2)
    plt.plot(x, mlu_averages, marker='o', label='MLU Average (%)')
    plt.title('MLU Average Over Time')
    plt.xlabel('Index')
    plt.ylabel('MLU Average (%)')
    plt.grid(True)

    plt.subplot(3, 1, 3)
    plt.plot(x, usages, marker='o', label='Usage (W)')
    plt.title('Power Usage Over Time')
    plt.xlabel('Index')
    plt.ylabel('Power Usage (W)')
    plt.grid(True)

    # 调整布局以避免重叠
    plt.tight_layout()

    # 显示图例
    plt.legend()

    # 保存图形为PNG图片
    plt.savefig(output_image_path)

    # 注意：这里不调用plt.show()，因为我们只保存图片而不显示它

if __name__ == "__main__":
    # 检查命令行参数数量
    if len(sys.argv) < 2 or len(sys.argv) > 3:
        print("Usage: python script.py <log_file_path> [<output_image_path>]")
        print("If <output_image_path> is not provided, the default 'default_system_metrics.png' will be used.")
        sys.exit(1)

    # 获取命令行参数
    log_file_path = sys.argv[1]
    output_image_path = sys.argv[2] if len(sys.argv) == 3 else 'default_system_metrics.png'

    # 调用函数绘制并保存图形
    plot_system_metrics(log_file_path, output_image_path)

